package com.glodon.paas.account.dao.redis;

import java.util.ArrayList;
import java.util.Collection;
import java.util.Collections;
import java.util.Set;

import org.springframework.data.redis.core.RedisTemplate;
import org.springframework.data.redis.core.ValueOperations;
import org.springframework.data.redis.serializer.JdkSerializationRedisSerializer;
import org.springframework.security.oauth2.common.OAuth2AccessToken;
import org.springframework.security.oauth2.common.OAuth2RefreshToken;
import org.springframework.security.oauth2.provider.OAuth2Authentication;
import org.springframework.security.oauth2.provider.token.AuthenticationKeyGenerator;
import org.springframework.security.oauth2.provider.token.DefaultAuthenticationKeyGenerator;
import org.springframework.security.oauth2.provider.token.TokenStore;

import redis.clients.jedis.Jedis;
import redis.clients.jedis.JedisPool;
import redis.clients.jedis.Transaction;

public class RedisTokenStore implements TokenStore {
	public static final String DEFAULT_ACCESS_PREFIX = "oauth2:accessToken:"; 	
	public static final String DEFAULT_REFRESH_PREFIX = "oauth2:refreshToken:";
	public static final String DEFAULT_AUTHKEY_PREFIX = "oauth2:authentication:";
	public static final String DEFAULT_NAME2ACCESS_PREFIX = "oauth2:usernameToAccess:";
	public static final String DEFAULT_ID2ACCESS_PREFIX = "oauth2:clientidToAccess:";	
	
	public static final String ACCESS_FIELD = "accessToken";
	public static final String REFRESH_FIELD = "refreshToken";
	public static final String AUTH_FIELD = "authentication";
	public static final String AUTH_KEY_FIELD = "authenticationKey";
	
	private final RedisTemplate<String, OAuth2AccessToken> accessTokenTemplate;
	private final ValueOperations<String, OAuth2AccessToken> accessTokenOperations;
	
	private final RedisTemplate<String, OAuth2RefreshToken> refreshTokenTemplate;
	private final ValueOperations<String, OAuth2RefreshToken> refreshTokenOperations;
	
	
	
	private JedisPool pool;
	
	private AuthenticationKeyGenerator authenticationKeyGenerator = new DefaultAuthenticationKeyGenerator();
	
	public RedisTokenStore(RedisTemplate<String, OAuth2AccessToken> accessTemplate,RedisTemplate<String, OAuth2RefreshToken> refreshTemplate){
		accessTokenTemplate = accessTemplate;
		accessTokenOperations = accessTokenTemplate.opsForValue();
		refreshTokenTemplate = refreshTemplate;
		refreshTokenOperations = refreshTokenTemplate.opsForValue();
	}
	
	public void setPool(JedisPool pool) {
		this.pool = pool;
	}
	
	@Override
	public OAuth2Authentication readAuthentication(OAuth2AccessToken token) {
		
		return readAuthentication(token.getValue());
	}

	@Override
	public OAuth2Authentication readAuthentication(String token) {
		final Jedis jedis = pool.getResource();
		try{			
			byte[] bytes = jedis.hget((DEFAULT_ACCESS_PREFIX+token).getBytes(), AUTH_FIELD.getBytes());
			return (OAuth2Authentication) new JdkSerializationRedisSerializer().deserialize(bytes);
			
		}finally{
			pool.returnResource(jedis);
		}	
	}

	@Override
	public void storeAccessToken(OAuth2AccessToken token, OAuth2Authentication authentication) {
		final Jedis jedis = pool.getResource();
		try {				
			Transaction transaction = jedis.multi();
			
			transaction.hset((DEFAULT_ACCESS_PREFIX + token.getValue()).getBytes(), ACCESS_FIELD.getBytes(), 
					new JdkSerializationRedisSerializer().serialize(token));
			
			transaction.hset((DEFAULT_ACCESS_PREFIX + token.getValue()).getBytes(), AUTH_FIELD.getBytes(), 
					new JdkSerializationRedisSerializer().serialize(authentication));
			//access token value to token
			//transaction.set((DEFAULT_ACCESS_PREFIX + token.getValue()).getBytes(),	new JdkSerializationRedisSerializer().serialize(token));	
			//access token value to authentication
			//transaction.set((DEFAULT_AUTH_PREFIX + token.getValue()).getBytes(),  new JdkSerializationRedisSerializer().serialize(authentication));
			
			final String key;
			key  = authenticationKeyGenerator.extractKey(authentication);
			//authentication key to access token
			transaction.set((DEFAULT_AUTHKEY_PREFIX+key).getBytes(), new JdkSerializationRedisSerializer().serialize(token));

			//user name to access token OR client id to access token
			if (!authentication.isClientOnly()) {
				transaction.sadd(DEFAULT_NAME2ACCESS_PREFIX+authentication.getName(), token.getValue());				
			}
			transaction.sadd(DEFAULT_ID2ACCESS_PREFIX+authentication.getAuthorizationRequest().getClientId(),token.getValue());			
		
			if (token.getExpiration() != null) {
				//TODO:set token TTL
				token.getValue();
				token.getExpiration();			
			}
			
			if (token.getRefreshToken() != null && token.getRefreshToken().getValue() != null) {
				
				transaction.hset((DEFAULT_REFRESH_PREFIX + token.getRefreshToken().getValue()).getBytes(), ACCESS_FIELD.getBytes(), 
						token.getValue().getBytes());
				
				//refresh token to access token
				//transaction.set(DEFAULT_REFRESH_PREFIX+token.getRefreshToken().getValue(), token.getValue());
				//access token to refresh token
				//be care of here 
				//transaction.set(token.getValue(), token.getRefreshToken().getValue());
			}
			
			transaction.exec();

		} finally {
			pool.returnResource(jedis);
		}	

	}

	@Override
	public OAuth2AccessToken readAccessToken(String tokenValue) {		
		//return accessTokenOperations.get(DEFAULT_ACCESS_PREFIX+tokenValue);
		
		final Jedis jedis = pool.getResource();
		try{			
			byte[] bytes = jedis.hget((DEFAULT_ACCESS_PREFIX+tokenValue).getBytes(), ACCESS_FIELD.getBytes());
			return (OAuth2AccessToken) new JdkSerializationRedisSerializer().deserialize(bytes);
			
		}finally{
			pool.returnResource(jedis);
		}		
		
		
	}

	@Override
	public void removeAccessToken(OAuth2AccessToken token) {
		removeAccessToken(token.getValue());
	}
	
	public void removeAccessToken(String tokenValue){
		final Jedis jedis = pool.getResource();
		try {				
			Transaction transaction = jedis.multi();
			OAuth2Authentication authentication = this.readAuthentication(tokenValue);
			
			if (authentication != null) {				
				transaction.srem(DEFAULT_NAME2ACCESS_PREFIX+authentication.getName(), tokenValue);
				
				transaction.srem(DEFAULT_ID2ACCESS_PREFIX+authentication.getAuthorizationRequest().getClientId(),tokenValue);
				//del auth key
				transaction.del(DEFAULT_AUTHKEY_PREFIX+authenticationKeyGenerator.extractKey(authentication));
				
			}		
			
			transaction.del((DEFAULT_ACCESS_PREFIX + tokenValue).getBytes());
			
			transaction.exec();
		}finally{
			pool.returnResource(jedis);
		}	
	}

	@Override
	public void storeRefreshToken(OAuth2RefreshToken refreshToken,
			OAuth2Authentication authentication) {
		final Jedis jedis = pool.getResource();
		try {				
			Transaction transaction = jedis.multi();
			
			transaction.hset((DEFAULT_REFRESH_PREFIX + refreshToken.getValue()).getBytes(), REFRESH_FIELD.getBytes(), 
					new JdkSerializationRedisSerializer().serialize(refreshToken));
			
			transaction.hset((DEFAULT_REFRESH_PREFIX + refreshToken.getValue()).getBytes(), AUTH_FIELD.getBytes(), 
					new JdkSerializationRedisSerializer().serialize(authentication));
			
			transaction.exec();
		}finally{
			pool.returnResource(jedis);
		}

	}

	@Override
	public OAuth2RefreshToken readRefreshToken(String tokenValue) {
		//return refreshTokenOperations.get(DEFAULT_REFRESH_PREFIX+tokenValue);
		final Jedis jedis = pool.getResource();
		try{			
			byte[] bytes = jedis.hget((DEFAULT_REFRESH_PREFIX+tokenValue).getBytes(), REFRESH_FIELD.getBytes());
			return (OAuth2RefreshToken) new JdkSerializationRedisSerializer().deserialize(bytes);
			
		}finally{
			pool.returnResource(jedis);
		}
	}

	@Override
	public OAuth2Authentication readAuthenticationForRefreshToken(
			OAuth2RefreshToken token) {
		final Jedis jedis = pool.getResource();
		try{			
			byte[] bytes = jedis.hget((DEFAULT_REFRESH_PREFIX+token.getValue()).getBytes(), AUTH_FIELD.getBytes());
			return (OAuth2Authentication) new JdkSerializationRedisSerializer().deserialize(bytes);
			
		}finally{
			pool.returnResource(jedis);
		}	
	}

	@Override
	public void removeRefreshToken(OAuth2RefreshToken token) {
		final Jedis jedis = pool.getResource();
		try {		
			jedis.del((DEFAULT_REFRESH_PREFIX + token.getValue()).getBytes());
						
		}finally{
			pool.returnResource(jedis);
		}	

	}

	@Override
	public void removeAccessTokenUsingRefreshToken(
			OAuth2RefreshToken refreshToken) {
		final Jedis jedis = pool.getResource();
		try {					
			String accessTokenValue = jedis.hget(DEFAULT_REFRESH_PREFIX + refreshToken.getValue(), ACCESS_FIELD).toString();
			if(null != accessTokenValue){
				this.removeAccessToken(accessTokenValue);
				jedis.hdel((DEFAULT_REFRESH_PREFIX + refreshToken.getValue()).getBytes(), ACCESS_FIELD.getBytes());				
			}			
						
		}finally{
			pool.returnResource(jedis);
		}

	}
	
	

	@Override
	public OAuth2AccessToken getAccessToken(OAuth2Authentication authentication) {
		final String key;
		key  = authenticationKeyGenerator.extractKey(authentication);
		final Jedis jedis = pool.getResource();
		try{			
			byte[] bytes = jedis.get((DEFAULT_AUTHKEY_PREFIX+key).getBytes());
			return (OAuth2AccessToken) new JdkSerializationRedisSerializer().deserialize(bytes);
			
		}finally{
			pool.returnResource(jedis);
		}	
		
	}

	@Override
	public Collection<OAuth2AccessToken> findTokensByUserName(String userName) {
		final Jedis jedis = pool.getResource();
		try{			
			Set<String> tokenSet = jedis.smembers(DEFAULT_NAME2ACCESS_PREFIX+userName);
			Collection<OAuth2AccessToken> tokenList = new ArrayList<OAuth2AccessToken>();
			for(String tokenKey: tokenSet)
			{
				OAuth2AccessToken entity = this.readAccessToken(tokenKey);
				tokenList.add(entity);			
			}	
			return tokenList != null ? Collections.<OAuth2AccessToken> unmodifiableCollection(tokenList) : Collections
					.<OAuth2AccessToken> emptySet();
			
		}finally{
			pool.returnResource(jedis);
		}	
	}

	@Override
	public Collection<OAuth2AccessToken> findTokensByClientId(String clientId) {
		final Jedis jedis = pool.getResource();
		try{			
			Set<String> tokenSet = jedis.smembers(DEFAULT_ID2ACCESS_PREFIX+clientId);
			Collection<OAuth2AccessToken> tokenList = new ArrayList<OAuth2AccessToken>();
			for(String tokenKey: tokenSet)
			{
				OAuth2AccessToken entity = this.readAccessToken(tokenKey);
				tokenList.add(entity);			
			}	
			return tokenList != null ? Collections.<OAuth2AccessToken> unmodifiableCollection(tokenList) : Collections
					.<OAuth2AccessToken> emptySet();
			
		}finally{
			pool.returnResource(jedis);
		}		
		

	}
	
	//just for IT test
	public OAuth2AccessToken getAccessTokenFromRefreshToken(OAuth2RefreshToken refreshToken){
		final Jedis jedis = pool.getResource();
		try {					
			if(jedis.hexists(DEFAULT_REFRESH_PREFIX + refreshToken.getValue(), ACCESS_FIELD)){
				String accessTokenValue = jedis.hget(DEFAULT_REFRESH_PREFIX + refreshToken.getValue(), ACCESS_FIELD).toString();
				if(null != accessTokenValue){				
					return this.readAccessToken(accessTokenValue);			
				}	
			}		
						
		}finally{
			pool.returnResource(jedis);
		}
		
		return null;
	}

}
